import rclpy                        # ROS2 Python接口库
from rclpy.node import Node         # ROS2 节点类
from sensor_msgs.msg import Image   # 图像消息类型
from cv_bridge import CvBridge      # ROS与OpenCV图像转换类
import cv2                          # Opencv图像处理库

OUTPUT_DIR = "output/"

class ImageRead(Node):
    def __init__(self, name):
        super().__init__(name)
        self.sub = self.create_subscription(Image, "image_capture", self.image_read_callback,3)
        self.transport = CvBridge()
        self.cnt = 0
    
    def image_read_callback(self, data):
        img = self.transport.imgmsg_to_cv2(data,"bgr8")
        self.get_logger().info(f"Recived an Image {self.cnt}")
        self.ImageProcess(img)

    def ImageProcess(self, img):
        for i in range(img.shape[0]):
            for j in range(img.shape[1]):
                R,G,B = img[i][j]
                avg = R/3 + G/3 + B/3
                img[i][j] = [avg,avg,avg]
        cv2.imwrite(f"{OUTPUT_DIR}output_{self.cnt}.jpg",img)
        self.cnt += 1

def main(arg=None):
    rclpy.init(args=arg)
    node = ImageRead("topic_image_sub")
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()
